
import numpy as np
import torch
import util.misc as utils
import itertools

@torch.no_grad()
def generate(model, post_processor, data_loader, device, verb_classes, missing_category_id):
    model.eval()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Generate:'

    detections = []
    preds = []
    for samples, targets in metric_logger.log_every(data_loader, 10, header):
        samples = samples.to(device)
        outputs = model(samples)
        orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0)
        results = post_processor(outputs, orig_target_sizes)
        #   每张图
        for img_results, img_targets in zip(results, targets):
            #每个交互
            for hoi in img_results['hoi_prediction']:
                detection = utils.CacheTemplate(image_id=img_targets['img_id'])
                detection['person_box'] = img_results['predictions'][hoi['subject_id']]['bbox'].tolist()
                if img_results['predictions'][hoi['object_id']]['category_id'] == missing_category_id:
                    object_box = [np.nan, np.nan, np.nan, np.nan]
                else:
                    object_box = img_results['predictions'][hoi['object_id']]['bbox'].tolist()        
                cut_agent = 0
                hit_agent = 0
                eat_agent = 0
                idx, score =hoi['category_id'], hoi['score']
                verb_class = verb_classes[idx]
                score = score.item()
                if len(verb_class.split('_')) == 1:
                    detection['{}_agent'.format(verb_class)] = score
                elif 'cut_' in verb_class:
                    detection[verb_class] = object_box + [score]
                    cut_agent = score if score > cut_agent else cut_agent
                elif 'hit_' in verb_class:
                    detection[verb_class] = object_box + [score]
                    hit_agent = score if score > hit_agent else hit_agent
                elif 'eat_' in verb_class:
                    detection[verb_class] = object_box + [score]
                    eat_agent = score if score > eat_agent else eat_agent
                else:
                    detection[verb_class] = object_box + [score]
                    detection['{}_agent'.format(verb_class.replace('_obj', '').replace('_instr', ''))] = score
                detection['cut_agent'] = cut_agent
                detection['hit_agent'] = hit_agent
                detection['eat_agent'] = eat_agent
                detections.append(detection)
            pass
        pass
    preds.extend(list(itertools.chain.from_iterable(utils.all_gather(detections))))        
    return preds
